import math
from collections import Counter   # Counter 用来统计邻居标签出现次数
from typing import List       # 为函数/变量添加类型提示，提高可读性和 IDE 支持

def euclidean_distance(a: List[float], b: List[float]) -> float:
    """
    计算两条等长向量的欧氏距离。
    """
    if len(a) != len(b):
        raise ValueError("向量维度不一致")# 防止维度不匹配导致 zip 失真
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))

class KNN:
    """
    手动实现的 K 近邻分类器。
    """

    def __init__(self, k: int, label_num: int):
        """
        :param k: 近邻数目
        :param label_num: 类别总数（用于统计）
        """
        if k <= 0:
            raise ValueError("k 必须为正整数")
        self.k = k
        self.label_num = label_num
        self.x_train: List[List[float]] = []   # 训练特征矩阵（列表的列表）
        self.y_train: List[int] = []          # 训练标签向量

    def fit(self, x_train: List[List[float]], y_train: List[int]) -> None:
        """
        保存训练数据。
        """
        if len(x_train) != len(y_train):
            raise ValueError("特征和标签数量不匹配")
        self.x_train = x_train
        self.y_train = y_train

    def _get_knn_indices(self, x: List[float]) -> List[int]:
        """
        返回距离样本 x 最近的 k 个训练样本的索引。
        """
        # 计算所有训练样本到 x 的距离
        distances = [euclidean_distance(sample, x) for sample in self.x_train]
        # 按距离升序排列，取前 k 个索引
        sorted_indices = sorted(range(len(distances)), key=lambda i: distances[i])
        return sorted_indices[:self.k]

    def _vote_label(self, knn_indices: List[int]) -> int:
        """
        根据 k 个最近邻的标签进行投票，返回出现次数最多的标签。
        """
        # 统计邻居标签出现次数
        label_counter = Counter(self.y_train[i] for i in knn_indices)
        # 选取出现次数最多的标签（若出现次数相同，取最小标签号）
        most_common_label, _ = max(label_counter.items(), key=lambda item: (item[1], -item[0]))
        return most_common_label

    def predict(self, x_test: List[List[float]]) -> List[int]:
        """
        对测试集进行预测，返回预测标签列表。
        """
        predictions = []
        for x in x_test:
            knn_idx = self._get_knn_indices(x)   # 找最近的 k 个邻居
            pred_label = self._vote_label(knn_idx)  # 投票决定标签
            predictions.append(pred_label)
        return predictions


# ------------------- 示例 -------------------
if __name__ == "__main__":
    # 简单的二维数据集
    X_train = [
        [1.0, 2.0],
        [1.5, 1.8],
        [5.0, 8.0],
        [6.0, 9.0],
        [1.0, 0.6],
        [9.0, 11.0],
    ]
    y_train = [0, 0, 1, 1, 0, 1]   # 两类标签 0 / 1

    X_test = [
        [2.0, 3.0],
        [8.0, 9.0],
        [0.5, 0.5],
    ]

    knn = KNN(k=3, label_num=2)   # k=3 表示看最近的 3 个邻居，label_num=2 表示共有 2 类
    knn.fit(X_train, y_train)    # 训练（其实只是保存数据）
    y_pred = knn.predict(X_test)  # 预测


    print("预测结果:", y_pred)   # 示例输出: [0, 1, 0]
